# coding:utf-8
# coding:utf-8
from timeseries.timeseries import TimeSeries
from statsmodels.tsa.arima_model import ARIMA
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

'''
ARIMA模型
'''


class Arima(TimeSeries):
    p = -1
    q = -1
    # 初始p,q
    _p = -1
    _q = -1

    def __init__(self):
        TimeSeries.__init__(self)
        self.algorithm_name = "ARIMA模型"
        self.ipynb_template_name = "arima-template.ipynb"
        
    def searchPAndQ(self, pmax=5, qmax=5):
        # 寻找最优的p和q
        bic_matrix = []
        for p in range(pmax + 1):
            tmp = []
            for q in range(qmax + 1):
                try:
                    tmp.append(ARIMA(self.inputs, (p, self.diff_degrees, q)).fit().bic)
                except:
                    tmp.append(None)
            bic_matrix.append(tmp)
        return pd.DataFrame(bic_matrix).stack().idxmin()
        
    def implent(self): 
        TimeSeries.implent(self)
        self.p = self._p
        self.q = self._q
        # 寻找最优的p及q值
        if self.p == -1 or self.q == -1:
            self.p, self.q = self.searchPAndQ()
        # 构造及训练模型
        self.algorithm = ARIMA(self.inputs, order=(self.p, self.diff_degrees, self.q)).fit()
        # 预测
        self.forcast_reuslt = self.algorithm.forecast(self.forcast_period_cnt)
        # 绘图
        self.drawChart()
        
    def drawChart(self):
        # 保存序列图
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(np.arange(0, len(self.inputs.values)), self.inputs.values, label="实际值")
        ax.plot(np.arange(len(self.inputs.values), len(self.inputs.values) + self.forcast_period_cnt), self.forcast_reuslt[0], label="预测值")
        ax.legend()
        ax.set_title("时序图")
        fig.savefig("%s/ts_result.png" % (self.chart_path))
        
    def prepareIpynbItems(self):
        TimeSeries.prepareIpynbItems(self)
        self.ipynb_items["#p#"] = self._p
        self.ipynb_items["#q#"] = self._q
    
    def saveToExcle(self):
        new_data_souce = self.data_source.copy()
        new_data_souce["_是否预测值"] = 0
        new_data_souce["_误差"] = 0
        time_values = new_data_souce[self.time_field_name].values
        minus = time_values[-1] - time_values[-2]
        for i in range(0, self.forcast_period_cnt):
            new_line = pd.DataFrame([[new_data_souce[self.time_field_name].values[-1] + minus, self.forcast_reuslt[0][i], 1, self.forcast_reuslt[1][i]]], columns=[self.time_field_name, self.value_field_name, "_是否预测值", "_误差"])
            new_data_souce = new_data_souce.append(new_line, ignore_index=True)
        new_data_souce.to_excel(self.result_excel_path)
